package master;

import java.awt.*;
import java.io.*;
import java.util.ArrayList;
import java.util.List;

import javax.swing.*;

import master.event.*;
import master.userInterface.*;


/**
 * A class for representing a quantum register
 * @author Luke Dobinson
 *
 */

public class QuantumRegister implements Serializable, RegisterEventInitiator, ComplexNumberListener{
	
	protected final ComplexNumber [] coeff;
	final int dimension;
	final int n;
	
	/**
	 * Constructor for specifying the number of qubits in the register
	 * @param n
	 */
	public QuantumRegister(int n){
		this.n = n;
		this.dimension = twoPow(n);
		this.coeff = new ComplexNumber[dimension];
		//MUST set to zero
		setZero();
		setToEigenstate(0);
		for(int i = 0 ; i < getDimension();i++)coeff[i].addComplexNumberListener(this);
	}
	
	/**
	 * Returns the bit string corresponding to a given eigenstate
	 * @param i
	 * @return
	 */
	public String bitString(int i){
		if(i<0 || i >= dimension) throw new IllegalArgumentException("Invalid number");
		String val = Integer.toBinaryString(i);
		while(val.length() < n){
			val = "0" + val;
		}
		return val;
	}
	
	/**
	 * return an array of bit strings corresponding to all possible eigenstates
	 * @return
	 */
	public String [] bitStrings(){
		String [] vals = new String[dimension];
		for(int i = 0 ; i < dimension ; i ++) vals[i] = bitString(i);
		return vals;
	}
	
	/**
	 * returns 2^n where n is a positive integer
	 * @param pow
	 * @return
	 */
	private int twoPow(int pow){
		int num = 1;
		if(pow < 0) throw new IllegalArgumentException("Must have a positive power");
		for(int j = 1 ; j <= pow ; j++) num = num * 2;
		return num;
	}
	
	/**
	 * provide a copy of one of the eigenstate's coefficients
	 * @return
	 */
	private ComplexNumber [] coeffCopy(){
		ComplexNumber [] newArr = new ComplexNumber[coeff.length];
		for(int i = 0 ; i < coeff.length ; i++) newArr[i] =  coeff[i].createNewInstance();
		return newArr;
	}
	
	/**
	 * set the value of all eigenstate coefficients to zero
	 */
	public void setZero() {
		for(int i = 0 ; i < coeff.length ; i++){
			if(coeff[i]==null) coeff[i] = new ComplexNumber(0,0);
			coeff[i].setReal(0);
			coeff[i].setImaginary(0);
		}
	}
	
	/**
	 * return a copy of one of the eigenstate coefficients
	 * @param n
	 * @return
	 */
	public ComplexNumber getCoefficient(int n) {
		if(n < 0 || n >= coeff.length)
			throw new IllegalArgumentException("invalid dimension");
		return coeff[n].createNewInstance();
	}
	
	/**
	 * return an array of copies of all coefficients
	 * @return
	 */
	public ComplexNumber[] getCoefficients() {
		return coeffCopy();
	}
	
	/**
	 * set the value of a coefficient
	 * @param b
	 * @param n
	 * @throws IllegalArgumentException
	 */
	public void setCoefficient(ComplexNumber b, int n) throws IllegalArgumentException {
		if(!(b instanceof ComplexNumber))
			throw new IllegalArgumentException("Please pass complex Number!");
		coeff[n].setReal(  b.getReal());
		coeff[n].setImaginary(  b.getImaginary());
		createRegisterEvent();
	}
	
	/**
	 * simultaniousl;y specifiy the values of all coefficients
	 * @param ComplexNumbers
	 * @throws IllegalArgumentException
	 */
	public void setCoefficients(ComplexNumber... ComplexNumbers)
			throws IllegalArgumentException {
		if(ComplexNumbers.length != coeff.length)
			throw new IllegalArgumentException("Please pass compatible array");
		for(int i = 0 ; i < coeff.length ; i++)
			setCoefficient(ComplexNumbers[i],i);
	}
	
	/**
	 * collapse the quantum register to the first eigenstate measured
	 */
	public void collapseToEigenstate() {
		setToEigenstate(nMeasure(1)[0]);
		normalise();
	}
	
	/**
	 * scale all coefficients so that the sum of their magnitudes squared is equal to one
	 */
	public void normalise() {
		float sumTot = (float)Math.sqrt(magSquared().getReal());
		
		if(Double.isNaN(sumTot)||Double.isInfinite(sumTot)||sumTot==0) sumTot = 1;
		
		//divide all elements by the sum total
		for(int i = 0 ; i < coeff.length ; i++){
			//System.out.print("Dividing " + coeff[i] + " mag " + coeff[i].magnitude().getReal() + " by "+sumTot);
			coeff[i].setReal(coeff[i].getReal()/sumTot);
			coeff[i].setImaginary(coeff[i].getImaginary()/sumTot);
			//System.out.println("\n\tmag = "+ coeff[i].magnitude());
		}
	}
	
	/**
	 * return the sum of the squares of the magnitudes of all coefficients, in an ideal 
	 * world this should equal to one!
	 * @return
	 */
	public ComplexNumber magSquared(){
		ComplexNumber tot = new ComplexNumber(0,0);
		for(int i = 0 ; i < getDimension();i++){
			//System.out.println(c + " * " + c.conjugate() + " = "+c.conjugate().multiply(c));
			tot.add(magSquared(i));
		}
		//System.out.println("Magnitude squared "+tot);
		
		return tot;
	}
	
	/**
	 * return this magnitude squared of a single coefficient
	 * @param eigenstate
	 * @return
	 */
	public ComplexNumber magSquared(int eigenstate){
		return getCoefficient(eigenstate).conjugate().multiply(getCoefficient(eigenstate));
	}
	
	/**
	 * manually set the register to a single eigenstate
	 * @param n
	 * @throws IllegalArgumentException
	 */
	public void setToEigenstate(int n) throws IllegalArgumentException {
		if(n <0 || n >= coeff.length)
			throw new IllegalArgumentException("no such eigenstate");
		float p = this.getCoefficient(n).phase().getReal();
		setZero();
		setCoefficient(ComplexNumber.expi(1, p), n);
		normalise();
	}
	
	/**
	 * return the dimension number of the state vector
	 * @return
	 */
	public int getDimension() {
		return dimension;
	}
	
	/**
	 * return the number of qubits making up the register
	 * @return
	 */
	public int getQBitLength() {
		return n;
	}
	
	/**
	 * convert the register to a string that can be read
	 */
	public String toString(){
		return n + " QBit Quantum Register (" + dimension + " possible states)";
	}
	
	/**
	 * draw each of the registers qubits
	 * @param g
	 * @param x
	 * @param y
	 * @param scale
	 */
	public void draw(Graphics g, int x, int y, float scale){
		for(Qubit b : getQubits())
			b.draw(g, x, y, scale);
	}
	
	/**
	 * get a panel containing the qubits, eigenstates, and button menus of the quantum register
	 * @return
	 */
	public JPanel QMDisplay(){
		/**
		 * Create The main display Panel
		 */
		JPanel mainPanel = new JPanel();
		
		
		/**
		 * divide the display up into left and right minimizable panels
		 */
		JPanel main = new JPanel();
		
		
		/**
		 * Now add the eigenstates readout
		 */
		String [] titles = new String[getDimension()];
		for(int i = 0 ; i < this.getDimension() ; i++)
			titles[i] = "|"+bitString(i)+">";
		JPanel eigenstates = ComplexNumber.display(200, false, false,titles, coeff);
		eigenstates.setSize(new Dimension(200,200*getDimension()));
		
		
		/**
		 * Add the bloch spheres display
		 */
		JPanel spheres = new JPanel();
		spheres.setLayout(new GridLayout(getQBitLength(),1));
		int dim = 200;
		for(int i = 0 ; i < getQBitLength() ; i++){
			JPanel p = new QubitPanel(main,getQubit(i),100,false);
			p.setPreferredSize(new Dimension(dim,dim));
			spheres.add(p);
		}
		
		/**
		 * Add the button menu
		 */
		JToolBar buttons = new QuantumRegisterButtonToolbar(this, main);
		
		main.setLayout(new GridBagLayout());
		
		GridBagConstraints c = new GridBagConstraints();
		
		c.fill = c.BOTH;
		
		c.gridx = 0;
		c.weightx = 1;
		c.weighty=0;
		main.add(buttons);
		
		c.weighty=0.9;
		c.gridy=1;
		main.add(
				new JSplitPane(JSplitPane.HORIZONTAL_SPLIT,
						new JScrollPane(spheres),
						new JScrollPane(eigenstates))
				, c);
		return main;
	}
	
	/**
	 * open the register in a new display frame
	 */
	public void display(String title){
		JFrame disp = new JFrame();
		disp.setTitle(this.toString());
		disp.add(QMDisplay());
		disp.setSize(600,800);
		disp.setVisible(true);
	}
	
	/**
	 * display a new frame with the string representation of the register as the title
	 */
	public void display(){
		display(this.toString());
	}
	
	/**
	 * parse a quantum register from a string for the form "qubits:" or "qubits:value1\nvalue2\netc"
	 * @param t
	 * @return
	 */
	public static QuantumRegister parseRegister(String t){
		//just a number
		String [] data = t.split(":");
		QuantumRegister q = new QuantumRegister(Integer.parseInt(data[0].replace("\n", "")));
		if(data.length>1){
			q.setCoefficients(ComplexNumber.parseArray(data[1]));
		}
		return q;
	}
	
	/**
	 * Encode the register in a parsable script format
	 */
	public String toScript(){
		String script = getQBitLength() + ":\n";
		boolean zero = true;
		for(ComplexNumber i : coeff)
			if(! i.isZero()) zero = false;
		if(zero) return script;
		
		for(ComplexNumber i : coeff)
			script += i.getReal() + ","+i.getImaginary()+"\n";
		return script;
	}
	
	/**
	 * return a qubit object corresponding to the specified qubit
	 * @param i
	 * @return
	 */
	public Qubit getQubit(int i){
		return new Qubit(this,i);
	}
	/**
	 * return an array of all qubitsd making up this register
	 * @return
	 */
	public Qubit [] getQubits(){
		Qubit [] qubits = new Qubit[this.getQBitLength()];
		for(int i = 0 ; i < this.getQBitLength(); i++) qubits[i] = getQubit(i);
		return qubits;
	}
	/**
	 * measure the specified qubits
	 * @param bits
	 */
	public void measureQBits(int ... bits){
		for(int b : bits)
			measureQBit(b);
	}
	/**
	 * measure a given qubit
	 * @param bit
	 */
	public void measureQBit(int bit){
		if(bit<0 || bit > n) throw new IllegalArgumentException("Invalid qubit");
		float zero = 0;//total sumSq of eigenstates with this qbit zero
		float one = 0;//total sumSq of eigenstates with this qbit one
		for(int i = 0 ; i < getDimension();i++){
			if( bitString(i).charAt(bit) == '1')
				one += magSquared(i).getReal();
			else if( bitString(i).charAt(bit) == '0')
				zero += magSquared(i).getReal();
		}
		//decide the fate of the qBit one + zero = 1 so
		if((one+zero)*Math.random()>one){
			//reduce the magnitude of the eigenstates where this bit is one to zero
			for(int i = 0 ; i < getDimension();i++)
				if( bitString(i).charAt(bit) == '1') setCoefficient(new ComplexNumber(0,0),i);
		}else{
			//reduce the magnitude of the eigenstates where this bit is zero to zero
			for(int i = 0 ; i < getDimension();i++)
				if( bitString(i).charAt(bit) == '0') setCoefficient(new ComplexNumber(0,0),i);
		}
		normalise();//return the eigenstates to normalisation
	}
	/**
	 * perform a real measurement n times. this, of course shouldn't be possible
	 * in the real world, but we can use it for debiugging etc
	 * @param n
	 * @return
	 */
	public int[] nMeasure(int n){
		if(n <= 0 ) throw new IllegalArgumentException("Invalid number of measurements");
		
		int [] results = new int[n];
		
		for(int step = 0 ; step < n ; step++){
		
			float sumTot = 0;
			//sum the magnitudes of the eigenstates squared
			for(ComplexNumber num : coeff){
				float m = ((ComplexNumber)num.conjugate().multiply(num)).getReal();
				sumTot += m;
			}
			
			float selectedMagnitude = (float) (Math.random() * sumTot);
			//System.out.println(selectedMagnitude+" / "+sumTot);
			sumTot = 0;
			//sum the magnitudes of the eigenstates squared
			for(int i = 0 ; i < coeff.length ; i++){
				//get the total magnitude up to this point
				float m = ((ComplexNumber)coeff[i].conjugate().multiply(coeff[i])).getReal();
				sumTot += m;
				//if the selected value is between them, collapse to eigenstate i
				if(sumTot >= selectedMagnitude){
					results[step] = i;
					i = coeff.length;//break the loop
				}else if(i == coeff.length-1){
					throw new IllegalArgumentException("Error with calculation");
				}
				
			}
		}
		return results;
	}
	
	/**
	 * return an entirely new quantum register that has its values copied from this one
	 * @return
	 */
	public QuantumRegister copy(){
		QuantumRegister ret = new QuantumRegister(getQBitLength());
		for(int i = 0 ; i < getDimension();i++){
			ret.setCoefficient(
					this.getCoefficient(i)
					, i);
		}
		return ret;
	}
	
	/**
	 * copy the values from another quantum register onto this one
	 * @param reg
	 */
	public void copyValues(QuantumRegister reg) throws IllegalArgumentException{
		if(this.getDimension() < reg.getDimension())
			throw new IllegalArgumentException("Registers incompatible");
		
		for(int i = 0 ; i < reg.getDimension();i++){
			setCoefficient(
					reg.getCoefficient(i)
					, i);
		}
	}

	
	List<RegisterEventListener> listeners = new ArrayList<RegisterEventListener>();
	
	public void addRegisterListener(RegisterEventListener l) {
		listeners.add(l);
	}

	
	public void createRegisterEvent() {
		RegisterEvent e = new RegisterEvent(this,0,this);
		for(RegisterEventListener l : listeners)
			l.registerEvent(e);
	}
	
	public void complexNumberValueChanged(ComplexNumberEvent e) {
		createRegisterEvent();
	}

	public void seeRegister(){
		for (int i = 0; i < dimension; i++){
			System.out.println("" + coeff[i]+ "|" + Integer.toBinaryString(i) + "("+ i +")>");
		}
	}
}